"""
模型训练和预测模块
包含所有机器学习模型的训练、保存和加载功能
"""

import pandas as pd
import numpy as np
import joblib
import streamlit as st
from sklearn.ensemble import RandomForestRegressor
from sklearn.svm import SVR
from sklearn.linear_model import ElasticNet
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score

from src.config import (
    MODELS_DIR, MODEL_FILES, RANDOM_STATE, TEST_SIZE, CV_FOLDS,
    RF_PARAM_GRID, SVM_PARAM_GRID, ELASTIC_PARAM_GRID
)
from src.data_loader import load_data


class ModelManager:
    """模型管理器，负责模型的训练、保存、加载和预测"""
    
    def __init__(self):
        self.models = {}
        self.scaler = None
        self.feature_names = None
        self.X_test = None
        self.y_test = None
        self.grid_searches = {}
        
    def train_models(self):
        """训练所有模型"""
        # 加载数据
        X, y, _ = load_data()
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE
        )
        
        # 保存测试集和特征名称
        self.X_test = X_test
        self.y_test = y_test
        self.feature_names = X.columns
        
        # 标准化数据
        self.scaler = StandardScaler()
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)
        
        # 训练随机森林
        self._train_random_forest(X_train, y_train)
        
        # 训练SVM
        self._train_svm(X_train_scaled, y_train)
        
        # 训练弹性网络
        self._train_elastic_net(X_train_scaled, y_train)
        
        # 保存模型
        self._save_models()
        
        return self.models, self.scaler, self.feature_names, X_test, y_test, self.grid_searches
    
    def _train_random_forest(self, X_train, y_train):
        """训练随机森林模型"""
        rf_grid = GridSearchCV(
            RandomForestRegressor(random_state=RANDOM_STATE, n_jobs=-1),
            RF_PARAM_GRID,
            cv=CV_FOLDS,
            scoring='r2',
            n_jobs=-1,
            verbose=0
        )
        rf_grid.fit(X_train, y_train)
        self.models['rf'] = rf_grid.best_estimator_
        self.grid_searches['rf'] = rf_grid
    
    def _train_svm(self, X_train_scaled, y_train):
        """训练SVM模型"""
        svm_grid = GridSearchCV(
            SVR(),
            SVM_PARAM_GRID,
            cv=CV_FOLDS,
            scoring='r2',
            n_jobs=-1,
            verbose=0
        )
        svm_grid.fit(X_train_scaled, y_train)
        self.models['svm'] = svm_grid.best_estimator_
        self.grid_searches['svm'] = svm_grid
    
    def _train_elastic_net(self, X_train_scaled, y_train):
        """训练弹性网络模型"""
        elastic_grid = GridSearchCV(
            ElasticNet(random_state=RANDOM_STATE),
            ELASTIC_PARAM_GRID,
            cv=CV_FOLDS,
            scoring='r2',
            n_jobs=-1,
            verbose=0
        )
        elastic_grid.fit(X_train_scaled, y_train)
        self.models['elastic'] = elastic_grid.best_estimator_
        self.grid_searches['elastic'] = elastic_grid
    
    def _save_models(self):
        """保存所有模型和相关信息"""
        MODELS_DIR.mkdir(exist_ok=True)
        
        # 保存模型
        joblib.dump(self.models['rf'], MODELS_DIR / MODEL_FILES['rf_model'])
        joblib.dump(self.models['svm'], MODELS_DIR / MODEL_FILES['svm_model'])
        joblib.dump(self.models['elastic'], MODELS_DIR / MODEL_FILES['elastic_model'])
        joblib.dump(self.scaler, MODELS_DIR / MODEL_FILES['scaler'])
        joblib.dump(self.feature_names, MODELS_DIR / MODEL_FILES['feature_names'])
        joblib.dump(self.X_test, MODELS_DIR / MODEL_FILES['X_test'])
        joblib.dump(self.y_test, MODELS_DIR / MODEL_FILES['y_test'])
        
        # 保存网格搜索对象
        joblib.dump(self.grid_searches['rf'], MODELS_DIR / MODEL_FILES['rf_grid'])
        joblib.dump(self.grid_searches['svm'], MODELS_DIR / MODEL_FILES['svm_grid'])
        joblib.dump(self.grid_searches['elastic'], MODELS_DIR / MODEL_FILES['elastic_grid'])
    
    def load_models(self):
        """加载已保存的模型"""
        try:
            # 加载模型
            self.models['rf'] = joblib.load(MODELS_DIR / MODEL_FILES['rf_model'])
            self.models['svm'] = joblib.load(MODELS_DIR / MODEL_FILES['svm_model'])
            self.models['elastic'] = joblib.load(MODELS_DIR / MODEL_FILES['elastic_model'])
            self.scaler = joblib.load(MODELS_DIR / MODEL_FILES['scaler'])
            self.feature_names = joblib.load(MODELS_DIR / MODEL_FILES['feature_names'])
            self.X_test = joblib.load(MODELS_DIR / MODEL_FILES['X_test'])
            self.y_test = joblib.load(MODELS_DIR / MODEL_FILES['y_test'])
            
            # 加载网格搜索对象
            self.grid_searches['rf'] = joblib.load(MODELS_DIR / MODEL_FILES['rf_grid'])
            self.grid_searches['svm'] = joblib.load(MODELS_DIR / MODEL_FILES['svm_grid'])
            self.grid_searches['elastic'] = joblib.load(MODELS_DIR / MODEL_FILES['elastic_grid'])
            
            return True
        except Exception as e:
            st.warning(f"⚠️ 模型加载失败: {e}")
            return False
    
    def predict(self, input_data, model_name):
        """
        使用指定模型进行预测
        
        Args:
            input_data (pd.DataFrame): 输入数据
            model_name (str): 模型名称 ('rf', 'svm', 'elastic')
            
        Returns:
            float: 预测结果
        """
        if model_name not in self.models:
            raise ValueError(f"模型 {model_name} 不存在")
        
        # 准备输入数据
        full_input = pd.DataFrame(columns=self.feature_names).fillna(0)
        for feature in input_data.columns:
            if feature in self.feature_names:
                full_input[feature] = input_data[feature]
        
        full_input = full_input.fillna(0)
        
        # 根据模型类型进行预测
        if model_name == 'rf':
            return self.models['rf'].predict(full_input)[0]
        else:
            # SVM和弹性网络需要标准化
            full_input_scaled = self.scaler.transform(full_input)
            return self.models[model_name].predict(full_input_scaled)[0]
    
    def get_model_performance(self, model_name):
        """
        获取模型性能指标
        
        Args:
            model_name (str): 模型名称
            
        Returns:
            dict: 性能指标字典
        """
        if model_name not in self.models:
            return None
        
        if model_name == 'rf':
            y_pred = self.models['rf'].predict(self.X_test)
        else:
            y_pred = self.models[model_name].predict(self.scaler.transform(self.X_test))
        
        mse = mean_squared_error(self.y_test, y_pred)
        r2 = r2_score(self.y_test, y_pred)
        
        return {
            'mse': mse,
            'r2': r2,
            'rmse': np.sqrt(mse)
        }
    
    def get_feature_importance(self):
        """获取随机森林的特征重要性"""
        if 'rf' not in self.models:
            return None
        
        importance_df = pd.DataFrame({
            'Feature': self.feature_names,
            'Importance': self.models['rf'].feature_importances_
        }).sort_values('Importance', ascending=False)
        
        return importance_df


@st.cache_resource
def get_model_manager():
    """获取模型管理器实例（带缓存）"""
    return ModelManager()


def load_or_train_models():
    """加载或训练模型"""
    model_manager = get_model_manager()
    
    # 检查模型文件是否存在
    all_files_exist = all(
        (MODELS_DIR / file).exists() 
        for file in MODEL_FILES.values()
    )
    
    if all_files_exist:
        st.info("📁 正在加载已保存的模型...")
        if model_manager.load_models():
            st.success("✅ 模型加载成功！")
            return model_manager
        else:
            st.info("🔄 将重新训练模型...")
    
    # 重新训练模型
    st.info("🚀 正在训练模型，这可能需要几分钟...")
    model_manager.train_models()
    return model_manager
